from io import BytesIO

import requests
from langchain_core.output_parsers import JsonOutputParser
from pydantic import BaseModel,Field
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph.message import add_messages
from langchain_core.messages import HumanMessage, AIMessage
from typing import TypedDict, Optional, List
#from langchain_core.pydantic_v1 import BaseModel, Field
#from matplotlib import image as mpimg, pyplot as plt


# 定义状态结构
class State(TypedDict):
    name: Optional[str]
    age: Optional[str]
    messages: List[HumanMessage | AIMessage]
    next: Optional[str]  # 流程控制字段
    user_initial_input: Optional[str]  # 存储初始用户输入


# 句子类型判断模型
class SentenceTypeJudgment(BaseModel):
    type: str = Field(
        description="句子类型，'declarative'表示陈述句，'interrogative'表示疑问句"
    )


# 信息提取模型
class UserInfoExtraction(BaseModel):
    name: Optional[str] = Field(None, description="提取的姓名")
    age: Optional[str] = Field(None, description="提取的年龄")


# 初始化模型
llm = ChatOpenAI(
    model_name="deepseek-chat",
    openai_api_base="https://api.deepseek.com/v1",
    openai_api_key="sk-f73f95fffd34497d8d9c11d868147710"
)

# 句子类型判断链
judge_prompt = ChatPromptTemplate.from_template("""
判断用户输入是陈述句还是疑问句：
1. 疑问句通常含疑问词（谁/什么/哪里）或末尾有问号，用于提问
2. 陈述句用于陈述事实，末尾通常为句号
仅返回JSON，type字段只能是'declarative'或'interrogative'

用户输入: {user_input}
{format_instructions}
""")
judge_parser = JsonOutputParser(pydantic_object=SentenceTypeJudgment)
judge_chain = judge_prompt | llm | judge_parser

# 信息提取链
extract_prompt = ChatPromptTemplate.from_template("""
从用户输入中提取姓名和年龄，未提及则返回None
仅返回JSON格式

用户输入: {user_input}
{format_instructions}
""")
extract_parser = JsonOutputParser(pydantic_object=UserInfoExtraction)
extract_chain = extract_prompt | llm | extract_parser


# 工具函数
def check_information_complete(state: State) -> bool:
    return state.get("name") is not None and state.get("age") is not None

def get_missing_fields(state: State) -> List[str]:
    missing = []
    if state.get("name") is None:
        missing.append("姓名")
    if state.get("age") is None:
        missing.append("年龄")
    return missing

def process_information(state: State) -> str:
    return f"已收集到您的信息：\n姓名：{state['name']}\n年龄：{state['age']}\n感谢提供！"


# 节点函数
def check_information_node(state: State) -> State:
    new_state = state.copy()
    if check_information_complete(state):
        result = process_information(state)
        new_state["messages"] = add_messages(state["messages"], [AIMessage(content=result)])
        new_state["next"] = END
    else:
        missing_text = "、".join(get_missing_fields(state))
        prompt = f"请补充以下信息：{missing_text}"
        new_state["messages"] = add_messages(state["messages"], [AIMessage(content=prompt)])
        new_state["next"] = "human_supplement"
    return new_state


# 修正：初始输入节点不直接获取输入，仅标记流程
def user_initial_input_node(state: State) -> State:
    new_state = state.copy()
    # 仅设置下一步为初始输入中断点，不在节点内调用input()
    new_state["next"] = "human_initial"
    return new_state


def judge_sentence_type_node(state: State) -> State:
    new_state = state.copy()
    user_input = new_state["user_initial_input"]

    try:
        # 判断句子类型
        judge_result = judge_chain.invoke({
            "user_input": user_input,
            "format_instructions": judge_parser.get_format_instructions()
        })
        sentence_type = judge_result["type"]

        if sentence_type == "declarative":
            # 疑问句：提取信息
            extract_result = extract_chain.invoke({
                "user_input": user_input,
                "format_instructions": extract_parser.get_format_instructions()
            })
            if extract_result.get("name"):
                new_state["name"] = extract_result["name"].strip()
            if extract_result.get("age"):
                new_state["age"] = extract_result["age"].strip()
            new_state["next"] = "check_information"

        elif sentence_type == "interrogative":
            # 陈述句：结束对话
            end_msg = "您输入的是一个问题，对话暂时结束。"
            new_state["messages"] = add_messages(state["messages"], [AIMessage(content=end_msg)])
            new_state["next"] = END

    except Exception as e:
        print(f"处理出错: {e}")
        new_state["messages"] = add_messages(state["messages"], [AIMessage(content="处理出错，对话结束。")])
        new_state["next"] = END

    return new_state


# 可视化函数
# def visualize_graph(graph):
#     try:
#         mermaid_code = graph.get_graph().draw_mermaid()
#         services = ["https://kroki.io/mermaid/png/"]
#         img_data = None
#         for service in services:
#             try:
#                 response = requests.post(
#                     service,
#                     data=mermaid_code.encode('utf-8'),
#                     headers={'Content-Type': 'text/plain'}
#                 )
#                 if response.status_code == 200:
#                     img_data = response.content
#                     break
#             except Exception as e:
#                 print(f"可视化服务出错: {e}")
#                 continue
#
#         if img_data:
#             img = mpimg.imread(BytesIO(img_data))
#             plt.figure(figsize=(12, 8))
#             plt.imshow(img)
#             plt.axis('off')
#             plt.title('对话流程图')
#             plt.tight_layout()
#             plt.show()
#         else:
#             print("无法生成流程图")
#
#     except Exception as e:
#         print(f"可视化出错: {e}")
#

# 输入处理函数
def process_initial_input(user_input: str, state: State) -> State:
    new_state = state.copy()
    new_state["user_initial_input"] = user_input
    new_state["messages"] = add_messages(state["messages"], [HumanMessage(content=user_input)])
    return new_state

def process_supplement_input(user_input: str, state: State) -> State:
    new_state = state.copy()
    try:
        extract_result = extract_chain.invoke({
            "user_input": user_input,
            "format_instructions": extract_parser.get_format_instructions()
        })
        if extract_result.get("name") and not new_state.get("name"):
            new_state["name"] = extract_result["name"].strip()
        if extract_result.get("age") and not new_state.get("age"):
            new_state["age"] = extract_result["age"].strip()
    except Exception as e:
        print(f"补充信息提取出错: {e}")
    new_state["messages"] = add_messages(state["messages"], [HumanMessage(content=user_input)])
    return new_state


# 创建工作流
def create_workflow():
    builder = StateGraph(State)

    # 添加节点
    builder.add_node("user_initial_input", user_initial_input_node)  # 初始输入引导节点
    builder.add_node("judge_sentence_type", judge_sentence_type_node)  # 句子判断节点
    builder.add_node("check_information", check_information_node)  # 信息检查节点
    builder.add_node("human_initial", lambda s: s)  # 初始输入中断节点
    builder.add_node("human_supplement", lambda s: s)  # 补充输入中断节点


    # 连接节点
    builder.add_edge(START, "user_initial_input")
    builder.add_edge("user_initial_input", "human_initial")  # 从引导节点到初始输入中断点
    builder.add_edge("human_initial", "judge_sentence_type")  # 输入后进入判断

    # 判断节点分支
    builder.add_conditional_edges(
        "judge_sentence_type",
        lambda x: x["next"],
        {"check_information": "check_information", END: END}
    )

    # 信息检查节点分支
    builder.add_conditional_edges(
        "check_information",
        lambda x: x["next"],
        {END: END, "human_supplement": "human_supplement"}
    )

    builder.add_edge("human_supplement", "check_information")

    # 编译工作流（指定所有中断点）
    memory = MemorySaver()
    graph = builder.compile(
        checkpointer=memory,
        interrupt_before=["human_initial", "human_supplement"]  # 所有需要用户输入的地方
    )
    #visualize_graph(graph)
    return graph


# 运行对话
def run_conversation():
    print("欢迎！请输入信息（例如：姓名张三，年龄25 / 请问需要提供姓名和年龄吗？）")

    workflow = create_workflow()
    thread_id = "user_thread_001"
    config = {"configurable": {"thread_id": thread_id}}

    # 初始化状态
    initial_state = {
        "name": None, "age": None, "messages": [],
        "next": None, "user_initial_input": None
    }
    workflow.update_state(config, initial_state)

    while True:
        # 运行流程直到中断点
        for event in workflow.stream(None, config, stream_mode="values"):
            current_state = event

        # 显示助手消息
        if current_state["messages"]:
            print(f"助手: {current_state['messages'][-1].content}")

        # 检查是否结束
        if current_state.get("next") == END:
            break

        # 处理初始输入
        if current_state["next"] == "human_initial":
            user_input = input("您: ")
            updated_state = process_initial_input(user_input, current_state)
            workflow.update_state(config, updated_state)

        # 处理补充输入
        elif current_state["next"] == "human_supplement":
            user_input = input("您: ")
            updated_state = process_supplement_input(user_input, current_state)
            workflow.update_state(config, updated_state)


if __name__ == "__main__":
    run_conversation()
